import {Dialog, Message, MessageDirection, MessageRole, MessageType, SessionConfig} from "@/types/chat";
import {GptVersion} from "@/app/constants";
import {create} from "zustand";
import {persist} from "zustand/middleware";
import {nanoid} from "nanoid";
import {completions} from "@/api";
import {useAccessStore} from './access';

interface ChatStore {
    id: number;
    sessions: ChatSession[];
    currentSessionIndex: number;
    openSession: (dialog?: { avatar?: string, title?: string }) => ChatSession;
    selectSession: (index: number) => void;
    deleteSession: (index: number) => void;
    currentSession: () => ChatSession;
    onSendMessage: (newMessage: Message) => Promise<void>;
    updateCurrentSession: (updater: (session: ChatSession) => void) => void;
    onRetry: () => void;
    deleteMessage: (message: Message) => void;
    createNewMessage: (value: string) => Message;
}

export interface ChatSession {
    // 会话ID
    id: number;
    // 对话框体
    dialog: Dialog;
    // 对话消息
    messages: Message[];
    // 会话配置
    config: SessionConfig;
    // 清除会话的索引
    clearContextIndex?: number;
}

function createChatSession(dialog?: { avatar?: string, title?: string }): ChatSession {
    return {
        id: 0,
        dialog: {
            avatar: dialog?.avatar ?? '/role/wali.png',
            title: dialog?.title ?? '新的对话',
            count: 0,
            subTitle: '请问有什么可以帮您？',
            timestamp: Date.now(),
        },
        messages: [{
            avatar: dialog?.avatar ?? '/role/wali.png',
            content: '请问有什么可以帮您？',
            messageType: MessageType.Text,
            time: Date.now(),
            direction: MessageDirection.Receive,
            role: MessageRole.system,
            id: nanoid()
        }],
        config: {
            gptVersion: GptVersion.GLM_4_AIR
        },
        clearContextIndex: undefined
    }
}

function formatMessages(messages: Message[]) {
    const latestMessages = messages.length > 3 ? messages.slice(-3) : messages;
    return latestMessages.map(({role, content}) => ({
        role, content
    }))
}

export function createNewMessage(value: string, role?: MessageRole) {
    return {
        avatar: role && MessageRole.user !== role ? "/role/wali.png" : "/role/runny-nose.png",
        content: value,
        messageType: MessageType.Text,
        time: Date.now(),
        role: role ?? MessageRole.user,
        id: nanoid(),
        streaming: false
    } as Message
}

export const userChatStore = create<ChatStore>()(
    persist(
        (set, get) => ({
            id: 0,
            sessions: [createChatSession()],
            currentSessionIndex: 0,
            openSession(dialog?: { avatar?: string, title?: string }) {
                const session = createChatSession(dialog);
                set(() => ({id: get().id + 1}));
                session.id = get().id;

                set((state) => ({
                    currentSessionIndex: 0,
                    sessions: [session].concat(state.sessions)
                }));
                return session;
            },
            selectSession(index: number) {
                set({
                    currentSessionIndex: index,
                });
            },
            deleteSession(index: number) {
                const count = get().sessions.length;
                const deleteSession = get().sessions.at(index);
                if (!deleteSession) return;

                const sessions = get().sessions.slice();
                sessions.splice(index, 1);

                const currentIndex = get().currentSessionIndex;
                let nextIndex = Math.min(
                    currentIndex - Number(index < currentIndex),
                    sessions.length - 1
                )
                if (count === 1) {
                    nextIndex = 0;
                    sessions.push(createChatSession());
                }

                set(() => ({
                    currentSessionIndex: nextIndex,
                    sessions
                }));
            },
            currentSession() {
                let index = get().currentSessionIndex;
                const sessions = get().sessions;
                if (index < 0 || index >= sessions.length) {
                    index = Math.min(sessions.length - 1, Math.max(0, index));
                    set(() => ({currentSessionIndex: index}));
                }
                return sessions[index];
            },
            async onSendMessage(newMessage: Message) {
                const session = get().currentSession();
                get().updateCurrentSession((session) => {
                    session.messages = session.messages.concat(newMessage);
                });

                const activeMessages = session.messages?.slice(session.clearContextIndex || 0);
                const messages = formatMessages(activeMessages);

                const botMessage = createNewMessage("", MessageRole.system);
                get().updateCurrentSession((session) => {
                    session.messages = session.messages.concat(botMessage);
                })

                const {body} = await completions({model: session.config.gptVersion, messages});
                const reader = body!.getReader();
                const decoder = new TextDecoder();

                new ReadableStream({
                    start(controller) {
                        async function push() {
                            const {done, value} = await reader.read();
                            if (done) {
                                controller.close();
                                return;
                            }
                            controller.enqueue(value);
                            const text = decoder.decode(value);

                            if (text === "0003") {
                                controller.close();
                                useAccessStore.getState().goToLogin();
                            }

                            botMessage.content += text;
                            get().updateCurrentSession((session) => {
                                session.messages = session.messages.concat();
                            });
                            push();
                        }

                        push();
                    }
                });


            },
            updateCurrentSession(updater) {
                const sessions = get().sessions;
                const index = get().currentSessionIndex;
                updater(sessions[index]);
                set(() => ({sessions}));
            },
            onRetry() {
                const session = get().currentSession();
                const activeMessages = session.messages.slice(session.clearContextIndex || 0);
                const messages = formatMessages(activeMessages);
                completions({model: session.config.gptVersion, messages});
            },
            deleteMessage(message: Message) {
                get().updateCurrentSession((session) => {
                    const index = session.messages.findIndex((m) => m.id === message.id);
                    session.messages.splice(index, 1);
                });
            },
            createNewMessage(value: string) {
                return {
                    avatar: "/role/runny-nose.png",
                    content: value,
                    time: Date.now(),
                    role: MessageRole.user,
                    id: nanoid(),
                } as Message;
            }
        }),
        {
            name: "chat-store"
        }
    )
)


